package com.sicau.rag;

import com.knuddels.jtokkit.api.EncodingType;
import org.springframework.ai.document.ContentFormatter;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.MetadataMode;
import org.springframework.ai.embedding.BatchingStrategy;
import org.springframework.ai.tokenizer.JTokkitTokenCountEstimator;
import org.springframework.ai.tokenizer.TokenCountEstimator;
import org.springframework.util.Assert;

import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

/**
 * 自定义批量策略
 */
public class TokenAndTrunkCountBatchingStrategy implements BatchingStrategy {
    private static final int DEFAULT_MAX_BATCH_SIZE = 10;
    private static final int MAX_INPUT_TOKEN_COUNT = 8191;
    private static final double DEFAULT_TOKEN_COUNT_RESERVE_PERCENTAGE = 0.1;
    private final int maxBatchSize;
    private final TokenCountEstimator tokenCountEstimator;
    private final int maxInputTokenCount;
    private final ContentFormatter contentFormatter;
    private final MetadataMode metadataMode;

    public TokenAndTrunkCountBatchingStrategy() {
        this(EncodingType.CL100K_BASE, 8191, 0.1, 10);
    }

    public TokenAndTrunkCountBatchingStrategy(EncodingType encodingType, int maxInputTokenCount, double reservePercentage, int maxBatchSize) {
        this(encodingType, maxInputTokenCount, reservePercentage, maxBatchSize, Document.DEFAULT_CONTENT_FORMATTER, MetadataMode.NONE);
    }

    public TokenAndTrunkCountBatchingStrategy(EncodingType encodingType, int maxInputTokenCount, double reservePercentage, int maxBatchSize, ContentFormatter contentFormatter, MetadataMode metadataMode) {
        Assert.notNull(encodingType, "EncodingType must not be null");
        Assert.isTrue(maxInputTokenCount > 0, "MaxInputTokenCount must be greater than 0");
        Assert.isTrue(reservePercentage >= (double)0.0F && reservePercentage < (double)1.0F, "ReservePercentage must be in range [0, 1)");
        Assert.notNull(contentFormatter, "ContentFormatter must not be null");
        Assert.notNull(metadataMode, "MetadataMode must not be null");
        Assert.isTrue(maxBatchSize > 0, "MaxBatchSize must be greater than 0");
        this.tokenCountEstimator = new JTokkitTokenCountEstimator(encodingType);
        this.maxInputTokenCount = (int)Math.round((double)maxInputTokenCount * ((double)1.0F - reservePercentage));
        this.maxBatchSize = maxBatchSize;
        this.contentFormatter = contentFormatter;
        this.metadataMode = metadataMode;
    }

    public List<List<Document>> batch(List<Document> documents) {
        List<List<Document>> batches = new ArrayList();
        int currentSize = 0;
        List<Document> currentBatch = new ArrayList();
        Map<Document, Integer> documentTokens = new LinkedHashMap();

        for(Document document : documents) {
            int tokenCount = this.tokenCountEstimator.estimate(document.getFormattedContent(this.contentFormatter, this.metadataMode));
            if (tokenCount > this.maxInputTokenCount) {
                throw new IllegalArgumentException("Tokens in a single document exceeds the maximum number of allowed input tokens");
            }

            documentTokens.put(document, tokenCount);
        }

        for(Document document : documentTokens.keySet()) {
            Integer tokenCount = (Integer)documentTokens.get(document);
            if (currentSize + tokenCount > this.maxInputTokenCount || currentBatch.size() >= this.maxBatchSize) {
                batches.add(currentBatch);
                currentBatch = new ArrayList();
                currentSize = 0;
            }

            currentBatch.add(document);
            currentSize += tokenCount;
        }

        if (!currentBatch.isEmpty()) {
            batches.add(currentBatch);
        }

        return batches;
    }
}

